package com.shujia.mllib

// 导入 Spark 和 MLlib 所需的类
import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.mllib.clustering.{KMeans, KMeansModel}
import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.sql.SparkSession

import scala.collection.mutable
// 导入 Spark ML 的评估器，用于计算聚类的 Silhouette 分数
import org.apache.spark.ml.evaluation.ClusteringEvaluator
object Test3 {
  def main(args: Array[String]): Unit = {
    // -------------------------
    // 1. 初始化 Spark
    // -------------------------
    // 创建 Spark 配置，设置应用名称和运行模式（local[*] 表示使用所有可用的本地 CPU 核心）
    val conf = new SparkConf()
      .setAppName("CovidClustering")
      .setMaster("local[*]")
    // 基于配置创建 SparkContext，用于底层 RDD 操作
    val sc = new SparkContext(conf)
    // 创建 SparkSession 对象，用于 DataFrame API 操作
    val spark = SparkSession.builder().getOrCreate()

    // 导入 SparkSession 内隐转换，方便 DataFrame 转换
    import spark.implicits._
    // -------------------------
    // 2. 加载数据并预处理
    // -------------------------
    // 通过 sc.textFile 加载 CSV 数据文件，注意路径应根据实际情况调整
    // 过滤掉第一行标题（假设标题行以 "temperature" 开头）
    val rawData = sc.textFile("spark/data/covid_data.csv")
      .filter(!_.startsWith("temperature"))
      .map { line =>
        // 以逗号分隔每行数据
        val fields = line.split(",")
        // 将前6个字段转换为 Double 类型并构建一个特征向量
        val features = Array(
          fields(0).toDouble, // 第1列：体温 (temperature)
          fields(1).toDouble, // 第2列：咳嗽 (cough)
          fields(2).toDouble, // 第3列：呼吸困难 (shortness_of_breath)
          fields(3).toDouble, // 第4列：旅行史 (travel_history)
          fields(4).toDouble, // 第5列：接触史 (contact_with_confirmed)
          fields(5).toDouble // 第6列：核酸检测 (pcr_test)
        )
        // 返回一个元组：(特征向量, 标签)，标签用于后续验证（但聚类本身是无监督学习）
        (Vectors.dense(features), fields(6))
      }
      .cache() // 缓存数据以提高后续计算效率
    // 将原始 RDD 数据转换为 DataFrame，方便后续使用 Spark ML 的评估器
    // 注意这里将 mllib 的 Vector 转换为 spark.ml 的 Vector
    val dataDF = rawData.map { case (features, label) =>
      (org.apache.spark.ml.linalg.Vectors.dense(features.toArray), label)
    }.toDF("features", "label")


    // -------------------------
    // 3. 数据划分：训练集与测试集
    // -------------------------
    // 使用 randomSplit 随机划分数据，比例 70% 作为训练集，30% 作为测试集，设置随机种子保证结果复现
    val Array(trainingDF, testDF) = dataDF.randomSplit(Array(0.8, 0.2), seed = 1L)
    trainingDF.cache() // 缓存训练集
    testDF.cache() // 缓存测试集
    // -------------------------
    // 4. 使用训练集训练 K-Means 模型
    // -------------------------
    // 由于使用的是 MLLib 的 KMeans.train 接口，需要将训练集 DataFrame 转换为 RDD[mllib.linalg.Vector]
    val trainingRDD = trainingDF.rdd.map(row =>
      // row.getAs[org.apache.spark.ml.linalg.Vector](0) 获取 "features" 列，再转换 为 mllib 的向量
        org.apache.spark.mllib.linalg.Vectors.dense(row.getAs[org.apache.spark.ml.linalg.Vector](0).toArray)
    ).cache()

    // 设置聚类参数：聚类数 k = 3，迭代次数为 100
    val numClusters = 3 // 假设将数据聚为 3 个簇（例如 Negative、Suspected、Confirmed）
    val numIterations = 100
    // 使用 KMeans.train 训练模型，返回 KMeansModel 对象
    val model: KMeansModel = KMeans.train(trainingRDD, numClusters,numIterations)
    // 输出聚类中心，每个中心代表一个簇的均值向量
    println("=== 聚类中心 ===")
    // 通过 model.clusterCenters 得到每个聚类的中心（质心），即每个簇内所有数据点的均值。
    // 利用 zipWithIndex 给每个中心分配一个编号。
    // 遍历所有聚类中心，并打印出每个簇的中心信息。
    // 这三个簇的中心是每个簇中所有数据点在各个特征上的平均值

    val clusterLabelMap = new mutable.HashMap[Int, String]()

    model.clusterCenters.zipWithIndex.foreach { case (center, idx) =>
      println(s"Cluster $idx center: ${center}")
      val t: Int = center.toArray(0).toInt
      if(t==35){
        clusterLabelMap.put(idx,"Negative")
      }else if(t==37){
        clusterLabelMap.put(idx,"Suspected")
      }else{
        clusterLabelMap.put(idx,"Confirmed")
      }
    }
    println("clusterLabelMap: "+clusterLabelMap)



    // 手动映射聚类编号到实际意义（此处仅作为参考，实际映射需结合数据具体情况）
//    val clusterLabelMap = Map(
//      0 -> "Confirmed", // 例如确诊
//      1 -> "Negative", // 例如低风险或阴性
//      2 -> "Suspected" // 例如疑似
//    )

    // -------------------------
    // 5. 对测试集进行预测与模型评估
    // -------------------------
    // 将测试集中的数据转换，并使用训练好的模型进行预测
    // 这里需要将测试集中的 features 从 spark.ml 向量转换为 mllib 向量，再进行预测
    val predictionsTest = testDF.rdd.map { row =>
      // 从 DataFrame 中获取特征向量和原始标签
      val features = row.getAs[org.apache.spark.ml.linalg.Vector](0)
      val label = row.getString(1)
      // 转换为 MLLib 的向量类型
      val mllibVector =  org.apache.spark.mllib.linalg.Vectors.dense(features.toArray)
      // 使用模型进行聚类预测，返回聚类编号
      val prediction = model.predict(mllibVector)
      // 将结果转换为 spark.ml 的向量类型，以便后续使用 Spark ML 的评估器
      (org.apache.spark.ml.linalg.Vectors.dense(features.toArray), label, prediction)
    }.toDF("features", "label", "prediction")
    // 展示测试集中部分预测结果
    println("=== 测试集部分预测结果 ===")
    predictionsTest.show(10, truncate = false)
    // 使用 ClusteringEvaluator 计算 Silhouette 分数来评估聚类效果
    // Silhouette 分数范围在 [-1,1]，越接近1表示聚类效果越好
    val evaluator = new ClusteringEvaluator()
      .setFeaturesCol("features") // 指定特征列
      .setPredictionCol("prediction") // 指定预测列
      .setMetricName("silhouette") // 指定评估指标为 silhouette
    val silhouette = evaluator.evaluate(predictionsTest)
    println(s"Silhouette score on test data = $silhouette")
    // -------------------------
    // 6. 用户交互预测
    // -------------------------
    println("=== 输入新患者症状进行聚类预测 ===")
    // 循环等待用户输入
    while (true) {
      print("请输入体温、咳嗽(0/1)、呼吸困难(0/1)、旅行史(0/1)、接触史(0/1)、核酸检测(0/1)（逗号分隔），或输入exit退出\n> ")
      val input = scala.io.StdIn.readLine().trim
      // 当输入 exit 时退出程序
      if (input.equalsIgnoreCase("exit")) System.exit(0)
      try {
        // 将输入的字符串按逗号分隔，并转换为 Double 数组
        val features = input.split(",").map(_.trim.toDouble)
        // 检查输入特征数量是否为6个
        require(features.length == 6, "必须输入6个数值特征")
        // 构造特征向量
        val vector = Vectors.dense(features)
        // 使用训练好的模型进行聚类预测
        val cluster = model.predict(vector)
        // 根据人工映射输出预测结果
        println(s"预测聚类结果：${clusterLabelMap(cluster)}")
      } catch {
        // 捕获输入错误并提示
        case e: Exception =>
          println(s"输入错误: ${e.getMessage}\n示例：38.5,1,0,1,0,1")
      }
    }
    // 关闭 SparkContext
    sc.stop()
  }
}